import numpy as np
import PIL.Image
import PIL.ImageDraw

__all__ = ['polygons_to_mask', 'keypoint_to_mask']

def polygons_to_mask(mask: np.ndarray, polygons):
    """
    draw mask by polygons
    :param img_shape:
    :param polygons:[(x,y),(x,y),.....]
    :return:
    """
    mask = PIL.Image.fromarray(mask)
    xy = list(map(tuple, polygons))
    PIL.ImageDraw.Draw(mask).polygon(xy=xy, outline=1, fill=1)
    mask = np.array(mask, dtype=bool)
    return mask


def keypoint_to_mask(landmark, image_shape):
    mask_point_left = []
    mask_point_right = []
    for i in range(int(landmark.shape[0] / 2)):
        l_p = landmark[i * 2]
        r_p = landmark[i * 2 + 1]
        mask_point_left.append([int(l_p[0]), int(l_p[1])])
        mask_point_right.append([int(r_p[0]), int(r_p[1])])
    mask_point = []
    mask_point.extend(mask_point_right)
    mask_point_left = mask_point_left[::-1]
    mask_point.extend(mask_point_left)
    mask = np.zeros(image_shape, dtype=np.uint8)
    mask = polygons_to_mask(mask, mask_point)
    mask = mask.astype(np.uint8)
    return mask




